"""
python fujian2_prophet_pridict_single_category.py
"""


import pandas as pd
import json
import os
from prophet import Prophet
import matplotlib.pyplot as plt

def Prophet_predict_single_category(input_json_path, output_json_path, graph_output_path):
    # 确保输出目录存在
    os.makedirs(os.path.dirname(output_json_path), exist_ok=True)
    os.makedirs(os.path.dirname(graph_output_path), exist_ok=True)

    # 读取JSON文件并转换为DataFrame
    with open(input_json_path, 'r') as file:
        data = json.load(file)
    df = pd.DataFrame(data)
    df['ds'] = pd.to_datetime(df['date'])
    df['y'] = df['sales']

    # 初始化Prophet模型并拟合数据
    model = Prophet()
    model.fit(df[['ds', 'y']])

    # 设置预测的时间范围为2023年7月1日到2023年9月30日
    future_dates = pd.date_range(start='2023-07-01', end='2023-09-30', freq='D')
    future = pd.DataFrame({'ds': future_dates})
    forecast = model.predict(future)

    # 格式化并输出预测结果
    forecast = forecast[['ds', 'yhat']].rename(columns={'ds': 'date', 'yhat': 'sales'})
    forecast['date'] = forecast['date'].dt.strftime('%Y-%m-%d')  # 将日期转换为字符串格式
    result_data = forecast.to_dict(orient='records')

    # 将预测数据写入输出JSON文件
    with open(output_json_path, 'w') as file:
        json.dump(result_data, file, indent=4)

    # 绘制原始数据和预测数据
    plt.figure(figsize=(10, 6))
    plt.plot(df['ds'], df['y'], label='Original Sales', color='blue')
    plt.plot(forecast['date'], forecast['sales'], label='Predicted Sales', color='orange')
    plt.xlabel('Date')
    plt.ylabel('Sales')
    plt.title('Sales Prediction')
    plt.legend()

    # 保存图表
    plt.savefig(graph_output_path)
    plt.close()

if __name__ == "__main__":
    # 示例调用
    input_path = "../fujian/fujian2/groupByCategory/category_category1.json"
    output_path = "../fujian/fujian2/prophet/json_output/category_category1_pred.json"
    graph_path = "../fujian/fujian2/prophet/graph/category_category1.png"
    Prophet_predict_single_category(input_path, output_path, graph_path)
